{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "67d5f991",
   "metadata": {},
   "outputs": [],
   "source": [
    "%run -i \"../util/file_utils.ipynb\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "099e6ea3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import csv\n",
    "import openai\n",
    "from llama_index import VectorStoreIndex\n",
    "from llama_index import Document\n",
    "openai.api_key = OPEN_AI_KEY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4b00c782",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('../data/IMDB-Movie-Data.csv') as f:\n",
    "    reader = csv.reader(f)\n",
    "    data = list(reader)\n",
    "    movies = data[1:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ea4e67a7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "id_='6e1ef633-f10b-44e3-9b77-f5f7b08dcedd' embedding=None metadata={'title': 'Guardians of the Galaxy', 'genres': ['Action', 'Adventure', 'Sci-Fi'], 'director': 'James Gunn', 'actors': ['Chris Pratt', ' Vin Diesel', ' Bradley Cooper', ' Zoe Saldana'], 'year': '2014', 'duration': '121', 'rating': '8.1', 'revenue': '333.13'} excluded_embed_metadata_keys=[] excluded_llm_metadata_keys=[] relationships={} hash='e18bdce3a36c69d8c1e55a7eb56f05162c68c97151cbaf4091814ae3df42dfe8' text='A group of intergalactic criminals are forced to work together to stop a fanatical warrior from taking control of the universe.' start_char_idx=None end_char_idx=None text_template='{metadata_str}\\n\\n{content}' metadata_template='{key}: {value}' metadata_seperator='\\n'\n",
      "id_='bf89acd7-fd92-46b0-88b2-eb163c5c04c6' embedding=None metadata={'title': 'Prometheus', 'genres': ['Adventure', 'Mystery', 'Sci-Fi'], 'director': 'Ridley Scott', 'actors': ['Noomi Rapace', ' Logan Marshall-Green', ' Michael Fassbender', ' Charlize Theron'], 'year': '2012', 'duration': '124', 'rating': '7', 'revenue': '126.46'} excluded_embed_metadata_keys=[] excluded_llm_metadata_keys=[] relationships={} hash='26a2f997159e1c3ab4d0a92a5b96903fb4e4ef9d5448731ba98fac104f70e373' text='Following clues to the origin of mankind, a team finds a structure on a distant moon, but they soon realize they are not alone.' start_char_idx=None end_char_idx=None text_template='{metadata_str}\\n\\n{content}' metadata_template='{key}: {value}' metadata_seperator='\\n'\n",
      "id_='80fe0547-ccbc-48ea-91e8-20161f9eba49' embedding=None metadata={'title': 'Split', 'genres': ['Horror', 'Thriller'], 'director': 'M. Night Shyamalan', 'actors': ['James McAvoy', ' Anya Taylor-Joy', ' Haley Lu Richardson', ' Jessica Sula'], 'year': '2016', 'duration': '117', 'rating': '7.3', 'revenue': '138.12'} excluded_embed_metadata_keys=[] excluded_llm_metadata_keys=[] relationships={} hash='664a7b57162852384e8ada2699f152f59c7c4ff4611eafbe23bcfbda565f5510' text='Three girls are kidnapped by a man with a diagnosed 23 distinct personalities. They must try to escape before the apparent emergence of a frightful new 24th.' start_char_idx=None end_char_idx=None text_template='{metadata_str}\\n\\n{content}' metadata_template='{key}: {value}' metadata_seperator='\\n'\n",
      "id_='1bd868a6-a85c-49aa-9372-3c1895ba3c3f' embedding=None metadata={'title': 'Sing', 'genres': ['Animation', 'Comedy', 'Family'], 'director': 'Christophe Lourdelet', 'actors': ['Matthew McConaughey', 'Reese Witherspoon', ' Seth MacFarlane', ' Scarlett Johansson'], 'year': '2016', 'duration': '108', 'rating': '7.2', 'revenue': '270.32'} excluded_embed_metadata_keys=[] excluded_llm_metadata_keys=[] relationships={} hash='4621637b94cdfe404a399590f6ce8c723310605beff8cd91b954a2de01571b2f' text=\"In a city of humanoid animals, a hustling theater impresario's attempt to save his theater with a singing competition becomes grander than he anticipates even as its finalists' find that their lives will never be the same.\" start_char_idx=None end_char_idx=None text_template='{metadata_str}\\n\\n{content}' metadata_template='{key}: {value}' metadata_seperator='\\n'\n",
      "id_='05c885b2-cc31-4a0e-b472-23bac0639167' embedding=None metadata={'title': 'Suicide Squad', 'genres': ['Action', 'Adventure', 'Fantasy'], 'director': 'David Ayer', 'actors': ['Will Smith', ' Jared Leto', ' Margot Robbie', ' Viola Davis'], 'year': '2016', 'duration': '123', 'rating': '6.2', 'revenue': '325.02'} excluded_embed_metadata_keys=[] excluded_llm_metadata_keys=[] relationships={} hash='e622080afd030d0a167a98e32b3a2ec6b370d0df410428d0e23dd3aaa304c7fb' text='A secret government agency recruits some of the most dangerous incarcerated super-villains to form a defensive task force. Their first mission: save the world from the apocalypse.' start_char_idx=None end_char_idx=None text_template='{metadata_str}\\n\\n{content}' metadata_template='{key}: {value}' metadata_seperator='\\n'\n",
      "id_='8976f868-795e-4803-9300-da4bef8d7444' embedding=None metadata={'title': 'The Great Wall', 'genres': ['Action', 'Adventure', 'Fantasy'], 'director': 'Yimou Zhang', 'actors': ['Matt Damon', ' Tian Jing', ' Willem Dafoe', ' Andy Lau'], 'year': '2016', 'duration': '103', 'rating': '6.1', 'revenue': '45.13'} excluded_embed_metadata_keys=[] excluded_llm_metadata_keys=[] relationships={} hash='777d8e68e41f2849fc1551970585bcbaad745c97d8b9ba779789d347aef4a531' text='European mercenaries searching for black powder become embroiled in the defense of the Great Wall of China against a horde of monstrous creatures.' start_char_idx=None end_char_idx=None text_template='{metadata_str}\\n\\n{content}' metadata_template='{key}: {value}' metadata_seperator='\\n'\n",
      "id_='d867832d-eca5-442e-a6ec-f594605d09e5' embedding=None metadata={'title': 'La La Land', 'genres': ['Comedy', 'Drama', 'Music'], 'director': 'Damien Chazelle', 'actors': ['Ryan Gosling', ' Emma Stone', ' Rosemarie DeWitt', ' J.K. Simmons'], 'year': '2016', 'duration': '128', 'rating': '8.3', 'revenue': '151.06'} excluded_embed_metadata_keys=[] excluded_llm_metadata_keys=[] relationships={} hash='6e24485f965c53202e36ab3410e015880d12c4b21aad8feddf115afd806cfcc3' text='A jazz pianist falls for an aspiring actress in Los Angeles.' start_char_idx=None end_char_idx=None text_template='{metadata_str}\\n\\n{content}' metadata_template='{key}: {value}' metadata_seperator='\\n'\n",
      "id_='6859e5ad-e167-4ea3-8dc2-38ce7da43688' embedding=None metadata={'title': 'Mindhorn', 'genres': ['Comedy'], 'director': 'Sean Foley', 'actors': ['Essie Davis', ' Andrea Riseborough', ' Julian Barratt', 'Kenneth Branagh'], 'year': '2016', 'duration': '89', 'rating': '6.4', 'revenue': ''} excluded_embed_metadata_keys=[] excluded_llm_metadata_keys=[] relationships={} hash='e74a993957f194b535281404e146d0bd3081b431666737c7d7ff14d228f6490c' text='A has-been actor best known for playing the title character in the 1980s detective series \"Mindhorn\" must work with the police when a serial killer says that he will only speak with Detective Mindhorn, whom he believes to be a real person.' start_char_idx=None end_char_idx=None text_template='{metadata_str}\\n\\n{content}' metadata_template='{key}: {value}' metadata_seperator='\\n'\n",
      "id_='10334b11-9c7e-46b9-beff-e487f6ed5721' embedding=None metadata={'title': 'The Lost City of Z', 'genres': ['Action', 'Adventure', 'Biography'], 'director': 'James Gray', 'actors': ['Charlie Hunnam', ' Robert Pattinson', ' Sienna Miller', ' Tom Holland'], 'year': '2016', 'duration': '141', 'rating': '7.1', 'revenue': '8.01'} excluded_embed_metadata_keys=[] excluded_llm_metadata_keys=[] relationships={} hash='f1929d3288716fc58fd363c661d7f266ef96dfed5817de89ab9db2b3594b3bde' text='A true-life drama, centering on British explorer Col. Percival Fawcett, who disappeared while searching for a mysterious city in the Amazon in the 1920s.' start_char_idx=None end_char_idx=None text_template='{metadata_str}\\n\\n{content}' metadata_template='{key}: {value}' metadata_seperator='\\n'\n",
      "id_='b9a7e732-b7df-435c-a31c-dd8d698754cd' embedding=None metadata={'title': 'Passengers', 'genres': ['Adventure', 'Drama', 'Romance'], 'director': 'Morten Tyldum', 'actors': ['Jennifer Lawrence', ' Chris Pratt', ' Michael Sheen', 'Laurence Fishburne'], 'year': '2016', 'duration': '116', 'rating': '7', 'revenue': '100.01'} excluded_embed_metadata_keys=[] excluded_llm_metadata_keys=[] relationships={} hash='2ea10a60973e97b5848d098ad2cc46ac3c0a4a20ebcb9aa5423424c9fb9cf6b9' text='A spacecraft traveling to a distant colony planet and transporting thousands of people has a malfunction in its sleep chambers. As a result, two passengers are awakened 90 years early.' start_char_idx=None end_char_idx=None text_template='{metadata_str}\\n\\n{content}' metadata_template='{key}: {value}' metadata_seperator='\\n'\n"
     ]
    }
   ],
   "source": [
    "documents = []\n",
    "for movie in movies[0:10]:\n",
    "    doc_id = movie[0]\n",
    "    title = movie[1]\n",
    "    genres = movie[2].split(\",\")\n",
    "    description = movie[3]\n",
    "    director = movie[4]\n",
    "    actors = movie[5].split(\",\")\n",
    "    year = movie[6]\n",
    "    duration = movie[7]\n",
    "    rating = movie[8]\n",
    "    revenue = movie[10]\n",
    "    document = Document(\n",
    "        text=description, \n",
    "        metadata={\n",
    "            \"title\": title, \n",
    "            \"genres\": genres, \n",
    "            \"director\": director, \n",
    "            \"actors\": actors, \n",
    "            \"year\": year, \n",
    "            \"duration\": duration, \n",
    "            \"rating\": rating, \n",
    "            \"revenue\": revenue\n",
    "        }\n",
    "    )\n",
    "    print(document)\n",
    "    documents.append(document)\n",
    "index = VectorStoreIndex.from_documents(documents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6af7a4e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine = index.as_query_engine()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2771b52d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "The Great Wall and Prometheus both talk about something gigantic. In The Great Wall, the protagonists become embroiled in the defense of the Great Wall of China against a horde of monstrous creatures. In Prometheus, the protagonists find a structure on a distant moon.\n"
     ]
    }
   ],
   "source": [
    "response = query_engine.query(\"\"\"Which movies talk about something gigantic?\"\"\")\n",
    "print(response.response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7af48399",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
